Feature: support rectilinear chunk grid extension #3534

jhamman · 2025-10-20T12:36:21Z

Summary

Adds support for RectilinearChunkGrid extension (Zarr v3), enabling arrays with variable chunk sizes per dimension.

Closes: #1595 | Replaces: #1483 | Related: zarr-extensions#25

Key Features

RectilinearChunkGrid

arr = zarr.create_array(
    shape=(60, 100),
    chunks=[[10, 20, 30], [25, 25, 25, 25]],
    zarr_format=3
)

Zarr v3 only (not compatible with v2, sharding, or from_array())
Supports RLE in JSON metadata: [[10, 6]] = 6 chunks of size 10
Stored internally in expanded format for fast indexing

Chunk Grid Access

grid = arr.chunk_grid                    # Returns ChunkGrid instance
grid.chunk_shapes                        # ((10, 20, 30), (25, 25, 25, 25))
isinstance(grid, RectilinearChunkGrid)   # Type-safe checking

`.chunks` Property Behavior

# RegularChunkGrid: returns tuple with FutureWarning (deprecated)
arr.chunks  # (10, 10)

# RectilinearChunkGrid: raises NotImplementedError
arr.chunks  # Use .chunk_grid instead

Design Decisions

Decision	Rationale
`ChunksLike` as TypeAlias	Flexible input types without runtime overhead
`ResolvedChunkSpec` as frozen dataclass	Named access, immutability, IDE support
Standalone validation functions	Testability, clear error messages, early validation
`.chunks` raises for rectilinear	No sensible single-tuple representation; guides users to `.chunk_grid`

Removed from Earlier Designs

Item	Reason
`RegularChunks`/`RectilinearChunks` tuple subclasses	Rejected - unnecessary complexity
Named dimension access (`chunks.lat`)	Removed per review feedback
`ChunksType` ABC hierarchy	Not implemented - TypeAlias approach preferred

Deferred / TODO Items

Item	Location	Notes
`update_shape()` optional chunks parameter	`metadata/v3.py:483`	Allow specifying new chunk sizes when resizing instead of default heuristic
Validation function placement	`chunk_grids.py:1513-1593`	Reviewer suggested moving to metadata module; kept for testability

Review Focus Areas

High Priority:

chunk_grids.py: RectilinearChunkGrid class, ChunksLike type, RLE expansion/compression, resolve_chunk_spec()
metadata/v3.py: update_shape() for rectilinear resize behavior
indexing.py: Variable chunk indexing with binary search
array.py: .chunks property behavior, .chunk_grid property

Tests:

test_chunk_grids/test_rectilinear.py: Comprehensive unit tests
test_chunk_grids/test_rectilinear_integration.py: End-to-end scenarios
testing/strategies.py: Hypothesis strategies for property-based testing

Breaking Changes

None. Fully backward compatible.

TODO:

Add unit tests and/or doctests in docstrings
Add docstrings and API docs for any new/modified user-facing classes and functions
New/modified features documented in docs/user-guide/*.md
Changes documented as a new file in changes/
GitHub Actions have all passed
Test coverage is 100% (Codecov passes)

…ature/rectilinear-chunk-grid

codecov · 2025-10-20T12:47:43Z

Codecov Report

❌ Patch coverage is 78.85533% with 133 lines in your changes missing coverage. Please review.
✅ Project coverage is 61.65%. Comparing base (b712f96) to head (41db2dc).

Files with missing lines	Patch %	Lines
src/zarr/core/chunk_grids.py	74.49%	76 Missing ⚠️
src/zarr/core/array.py	75.00%	20 Missing ⚠️
src/zarr/core/indexing.py	85.07%	20 Missing ⚠️
src/zarr/testing/strategies.py	89.28%	9 Missing ⚠️
src/zarr/core/metadata/v2.py	72.72%	6 Missing ⚠️
src/zarr/core/_info.py	0.00%	1 Missing ⚠️
src/zarr/core/metadata/v3.py	90.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3534      +/-   ##
==========================================
+ Coverage   60.94%   61.65%   +0.71%     
==========================================
  Files          86       86              
  Lines       10268    10769     +501     
==========================================
+ Hits         6258     6640     +382     
- Misses       4010     4129     +119

Files with missing lines	Coverage Δ
src/zarr/api/asynchronous.py	`72.20% <ø> (ø)`
src/zarr/api/synchronous.py	`36.61% <ø> (ø)`
src/zarr/core/group.py	`70.27% <ø> (ø)`
src/zarr/core/_info.py	`51.80% <0.00%> (ø)`
src/zarr/core/metadata/v3.py	`59.91% <90.00%> (+1.90%)`	⬆️
src/zarr/core/metadata/v2.py	`60.31% <72.72%> (+2.17%)`	⬆️
src/zarr/testing/strategies.py	`94.18% <89.28%> (-3.66%)`	⬇️
src/zarr/core/array.py	`67.99% <75.00%> (-0.12%)`	⬇️
src/zarr/core/indexing.py	`70.19% <85.07%> (+0.73%)`	⬆️
src/zarr/core/chunk_grids.py	`70.70% <74.49%> (+8.40%)`	⬆️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

…ature/rectilinear-chunk-grid

docs/user-guide/arrays.md

d-v-b · 2025-10-20T21:28:48Z

src/zarr/core/chunk_grids.py

+
+
+@dataclass(frozen=True)
+class RectilinearChunkGrid(ChunkGrid):


thoughts on just calling this class Rectilinear, and renaming the RegularChunkGrid to Regular? We could keep around a RegularChunkGrid class for compatibility. But I feel like people know these are chunk grids when they import them

50/50. I think the more descriptive class is useful when looking at a tracebacks. Plus, this is currently in .core so its not meant to be used directly by users.

src/zarr/core/chunk_grids.py

src/zarr/core/indexing.py

src/zarr/testing/strategies.py

dcherian · 2025-10-27T21:32:15Z

tests/test_properties.py

-@given(data=st.data())
-async def test_basic_indexing(data: st.DataObject) -> None:
-    zarray = data.draw(simple_arrays())
+@given(data=st.data(), zarray=st.one_of([simple_arrays(), complex_chunked_arrays()]))


Because the search space for the standard arrays strategy is so large, i made a different one complex_chunked_arrays that purely checks different chunk grids
with simple_arrays() we are only spending 10% of our time trying RectilinearChunkGrid so using this approach. We should boost number of examples too.

shoyer · 2025-10-28T03:11:15Z

docs/user-guide/arrays.md

+2. **Not compatible with sharding**: You cannot use variable chunking together with
+   the sharding feature. Arrays must use either variable chunking or sharding, but not both.


I hope this is a temporary limitation! There's a natural extension of rectilinear chunk grids to rectilinear shard grids.

xref: zarr-developers/zarr-extensions#34

Co-authored-by: Davis Bennett <[email protected]>

keewis · 2026-01-08T21:13:23Z

is it intentional that the chunk grid is RectilinearChunkGrid when explicitly enumerating regular chunk sizes?

import zarr

store = zarr.storage.MemoryStore()
arr = zarr.create_array(
    store,
    shape=(60, 100),
    chunks=[[20, 20, 20], [25, 25, 25, 25]],
    zarr_format=3,
    dtype="uint8"
)
arr.metadata.chunk_grid # RectilinearChunkGrid

Or is RegularChunkGrid only chosen if I specify a tuple of values, in this case (20, 25)?

Not sure how much this matters in practice (xarray's code does not appear to be affected by the difference), but that's what dask normalizes the chunks to.

jhamman · 2026-01-14T16:58:17Z

is it intentional that the chunk grid is RectilinearChunkGrid when explicitly enumerating regular chunk sizes?

@keewis - yes, this is intentional. We decided this was preferable because it will allow the user to add variable lengths later (through an extend/append workflow).

codspeed-hq · 2026-01-16T18:43:13Z

CodSpeed Performance Report

Merging this PR will not alter performance

_{Comparing jhamman:feature/rectilinear-chunk-grid (665847a) with main (e94504e)}

Summary

✅ 30 untouched benchmarks

src/zarr/core/chunk_grids.py

d-v-b · 2026-01-26T19:25:47Z

src/zarr/core/chunk_grids.py

+
+
+@dataclass(frozen=True)
+class ResolvedChunkSpec:


does this need to be a class? It has no methods. Seems like either a TypedDict or just a tuple would work

d-v-b · 2026-01-26T19:26:57Z

src/zarr/core/chunk_grids.py

+    shards: tuple[int, ...] | None
+
+
+def _validate_zarr_format_compatibility(


why do we need this here? Can't we check when we construct the array metadata document?

d-v-b · 2026-01-26T19:29:41Z

src/zarr/core/chunk_grids.py

+            )
+
+
+def _validate_sharding_compatibility(


shouldn't this be defined over in the array v3 metadata module, and used there to check that the array metadata document is valid?

d-v-b · 2026-01-26T19:31:26Z

src/zarr/core/chunk_grids.py

+            )
+
+
+def _validate_data_compatibility(


does this need to be stand-alone function? are we going to use this logic anywhere other than it's current usage site (resolve_chunk_spec)

src/zarr/core/metadata/v3.py

d-v-b · 2026-01-26T19:51:50Z

src/zarr/core/metadata/v3.py

@@ -436,6 +468,21 @@ def to_dict(self) -> dict[str, JSON]:
        return out_dict

    def update_shape(self, shape: tuple[int, ...]) -> Self:


IMO this needs to take a parameter that can define the new chunks (which can default to None or some other sentinel flagging "default behavior" semantics)

src/zarr/core/array.py

tomwhite · 2026-01-29T13:05:11Z

FYI I tested this PR for implementing rechunking with variable-sized intermediate chunks in Cubed - and it worked!

I found one wrinkle in that zarr.create_array supports rectilinear chunk grids, but zarr.open (with mode="w") doesn't. But that could be addressed later.

…arr-python into feature/rectilinear-chunk-grid

…ature/rectilinear-chunk-grid

d-v-b · 2026-02-04T07:10:15Z

I will give this a spin today, but assuming everything works, my question is: how can we ship this soon (this week?) while making it clear that the feature is experimental?

tinaok · 2026-02-04T17:07:06Z

I have varient chunked array;

import xarray as xr
from healpix_geo import nested 
import numpy as np

da = xr.open_zarr(
    "https://data-taos.ifremer.fr/GRID4EARTH/no_chunk_healpix.zarr",
    consolidated=True,   # if metadata is consolidated
).da

depth = da.cell_ids.attrs['level']
new_depth = depth-6
parents = nested.zoom_to(da.cell_ids, depth=depth, new_depth=new_depth) 
_, chunk_sizes =np.unique(parents, return_counts=True)
da.chunk({"cell_ids": tuple(chunk_sizes.tolist())})

Happy to experiment to write it using this experimental feature(if it works with xarray...)

jhamman added 5 commits October 16, 2025 13:32

feature(chunkgrids): add rectillinear chunk grid metadata support

0fd633f

implementation (wip)

31c831d

fixup

d40dd30

fixup types

5f84198

Merge branch 'main' of github.com:zarr-developers/zarr-python into fe…

b96a650

…ature/rectilinear-chunk-grid

github-actions bot added the needs release notes Automatically applied to PRs which haven't added release notes label Oct 20, 2025

jhamman added 3 commits October 20, 2025 09:42

more tests

6deb98e

Merge branch 'main' of github.com:zarr-developers/zarr-python into fe…

fc07af1

…ature/rectilinear-chunk-grid

docs

527f7cf

github-actions bot removed the needs release notes Automatically applied to PRs which haven't added release notes label Oct 20, 2025

fixups

cea5fc2

d-v-b reviewed Oct 20, 2025

View reviewed changes

docs/user-guide/arrays.md Show resolved Hide resolved

d-v-b reviewed Oct 20, 2025

View reviewed changes

docs/user-guide/arrays.md Outdated Show resolved Hide resolved

d-v-b reviewed Oct 20, 2025

View reviewed changes

src/zarr/core/chunk_grids.py Show resolved Hide resolved

d-v-b reviewed Oct 20, 2025

View reviewed changes

src/zarr/core/chunk_grids.py Show resolved Hide resolved

d-v-b reviewed Oct 20, 2025

View reviewed changes

src/zarr/core/indexing.py Show resolved Hide resolved

keewis mentioned this pull request Oct 27, 2025

variable-sized chunks with zarr v3 pydata/xarray#10880

Draft

3 tasks

dcherian added 2 commits October 27, 2025 14:34

Updated strategies

cef873f

better approach

79ecee9

dcherian reviewed Oct 27, 2025

View reviewed changes

src/zarr/testing/strategies.py Outdated Show resolved Hide resolved

dcherian reviewed Oct 27, 2025

View reviewed changes

shoyer reviewed Oct 28, 2025

View reviewed changes

jhamman and others added 6 commits October 28, 2025 09:57

handle run length encoded chunk grid spec in the top level api

e116f63

fixups

6eab392

fixups + support partial edge chunks / resize

1935d19

Update docs/user-guide/arrays.md

4f36ffb

Co-authored-by: Davis Bennett <[email protected]>

Merge branch 'main' into feature/rectilinear-chunk-grid

4c33eac

fixups + support partial edge chunks / resize

99fd0a8

Merge branch 'main' into feature/rectilinear-chunk-grid

9024817

doc notes about this being experimental until 3.3

84f47d3

Merge branch 'main' into feature/rectilinear-chunk-grid

665847a